#!/usr/bin/env python3
"""
Mock NodeManager 启动命令接收服务器
模拟 NodeManager 接收来自 Controller 的启动命令
"""

import json
import time
import threading
import asyncio
from typing import Dict, Any, List
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
import uvicorn
from pydantic import BaseModel, ValidationError

from motor.resources.http_msg_spec import StartCmdMsg, Ranktable, ServerInfo
from motor.resources.endpoint import Endpoint, DeviceInfo
from motor.resources.instance import ParallelConfig
from motor.utils.logger import get_logger

logger = get_logger(__name__)

# 全局变量存储接收到的启动命令
received_commands: List[Dict[str, Any]] = []
command_lock = threading.Lock()

app = FastAPI(title="Mock NodeManager Start Command Server", version="1.0.0")


class StartCommandResponse(BaseModel):
    """启动命令响应格式"""
    status: str
    message: str
    job_name: str
    instance_id: int
    endpoints_count: int
    ranktable_info: Dict[str, Any]


@app.get("/")
async def root():
    """根路径，返回服务器状态"""
    return {
        "service": "Mock NodeManager Start Command Server",
        "status": "running",
        "received_commands_count": len(received_commands)
    }


@app.get("/node-manager/status")
async def get_status():
    """获取服务器状态"""
    with command_lock:
        return {
            "status": "running",
            "received_commands": len(received_commands),
            "commands": received_commands[-5:] if received_commands else []  # 返回最近5个命令
        }


@app.post("/node-manager/start")
async def receive_start_command(request: Request):
    """
    接收启动命令
    
    这是 NodeManager 接收 Controller 发送的启动命令的端点
    """
    try:
        # 获取请求体
        body = await request.json()
        logger.info(f"收到启动命令: {json.dumps(body, indent=2, ensure_ascii=False)}")
        
        # 验证启动命令格式
        try:
            start_cmd = StartCmdMsg(**body)
            logger.info(f"启动命令验证成功: job_name={start_cmd.job_name}, "
                       f"instance_id={start_cmd.instance_id}, role={start_cmd.role}")
        except ValidationError as e:
            logger.error(f"启动命令格式验证失败: {e}")
            raise HTTPException(status_code=400, detail=f"Invalid start command format: {e}")
        
        # 存储接收到的命令
        command_data = {
            "timestamp": time.time(),
            "job_name": start_cmd.job_name,
            "role": start_cmd.role,
            "instance_id": start_cmd.instance_id,
            "endpoints_count": len(start_cmd.endpoints),
            "ranktable": start_cmd.ranktable.model_dump(),
            "raw_data": body
        }
        
        with command_lock:
            received_commands.append(command_data)
        
        # 模拟处理启动命令
        logger.info(f"开始处理启动命令: {start_cmd.job_name}")
        
        # 模拟启动推理引擎
        await simulate_engine_startup(start_cmd)
        
        # 返回成功响应
        response = StartCommandResponse(
            status="success",
            message=f"Start command received and processed for job {start_cmd.job_name}",
            job_name=start_cmd.job_name,
            instance_id=start_cmd.instance_id,
            endpoints_count=len(start_cmd.endpoints),
            ranktable_info={
                "version": start_cmd.ranktable.version,
                "status": start_cmd.ranktable.status,
                "server_count": start_cmd.ranktable.server_count,
                "server_list_count": len(start_cmd.ranktable.server_list)
            }
        )
        
        logger.info(f"启动命令处理完成: {start_cmd.job_name}")
        return response.model_dump()
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"处理启动命令时发生错误: {e}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")


async def simulate_engine_startup(start_cmd: StartCmdMsg):
    """
    模拟推理引擎启动过程
    
    Args:
        start_cmd: 启动命令
    """
    logger.info(f"模拟启动推理引擎: {start_cmd.job_name}")
    
    # 模拟启动延迟
    await asyncio.sleep(0.1)
    
    # 模拟启动步骤
    steps = [
        "初始化推理引擎",
        "加载模型",
        "配置并行参数",
        "启动推理服务",
        "注册到服务发现"
    ]
    
    for i, step in enumerate(steps, 1):
        logger.info(f"步骤 {i}/{len(steps)}: {step}")
        await asyncio.sleep(0.05)  # 模拟处理时间
    
    logger.info(f"推理引擎启动完成: {start_cmd.job_name}")


@app.get("/node-manager/commands")
async def get_all_commands():
    """获取所有接收到的启动命令"""
    with command_lock:
        return {
            "total_commands": len(received_commands),
            "commands": received_commands
        }


@app.delete("/node-manager/commands")
async def clear_commands():
    """清空所有接收到的启动命令"""
    with command_lock:
        received_commands.clear()
    return {"message": "All commands cleared"}


@app.get("/node-manager/commands/{job_name}")
async def get_commands_by_job(job_name: str):
    """根据作业名称获取启动命令"""
    with command_lock:
        job_commands = [cmd for cmd in received_commands if cmd["job_name"] == job_name]
        return {
            "job_name": job_name,
            "commands_count": len(job_commands),
            "commands": job_commands
        }


class MockNodeManagerServer:
    """Mock NodeManager 服务器类"""
    
    def __init__(self, host: str = "0.0.0.0", port: int = 8080):
        self.host = host
        self.port = port
        self.server = None
        self.server_thread = None
        
    def start(self):
        """启动服务器"""
        def run_server():
            config = uvicorn.Config(
                app,
                host=self.host,
                port=self.port,
                log_level="info"
            )
            self.server = uvicorn.Server(config)
            self.server.run()
        
        self.server_thread = threading.Thread(target=run_server, daemon=True)
        self.server_thread.start()
        logger.info(f"Mock NodeManager 服务器启动在 {self.host}:{self.port}")
        
    def stop(self):
        """停止服务器"""
        if self.server:
            self.server.should_exit = True
        logger.info("Mock NodeManager 服务器已停止")
    
    def get_received_commands(self):
        """获取接收到的命令"""
        with command_lock:
            return received_commands.copy()


def main():
    """主函数"""
    import asyncio
    
    print("=== Mock NodeManager 启动命令接收服务器 ===")
    print("这个服务器模拟 NodeManager 接收来自 Controller 的启动命令")
    print("启动后，Controller 会将启动命令发送到此服务器的 /start 端点")
    print()
    
    # 创建并启动服务器
    server = MockNodeManagerServer(host="0.0.0.0", port=9090)
    server.start()
    
    try:
        print(f"服务器运行在: http://localhost:{server.port}")
        print("可用的端点:")
        print(f"  GET  /                    - 服务器状态")
        print(f"  GET  /status              - 获取状态")
        print(f"  POST /start               - 接收启动命令")
        print(f"  GET  /commands            - 获取所有命令")
        print(f"  GET  /commands/{{job_name}} - 获取特定作业的命令")
        print(f"  DELETE /commands          - 清空所有命令")
        print()
        print("按 Ctrl+C 停止服务器")
        
        # 保持服务器运行
        while True:
            time.sleep(1)
            
    except KeyboardInterrupt:
        print("\n正在停止服务器...")
        server.stop()
        print("服务器已停止")


if __name__ == "__main__":
    main()
